package com.burt.zookeeper.javaapi.lock;

import java.util.List;
import java.util.Random;
import java.util.SortedSet;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;

/**
 * Description: 分布式锁(共享锁) User: Burt Date: 2017-08-11 Time: 12:06
 */
public class DistributeLock {

    private static final String ROOT_LOCKS = "/LOCKS";//根节点

    private final static byte[] data = {1, 2}; //节点的数据

    private CountDownLatch countDownLatch = new CountDownLatch(1);

    private ZooKeeper zooKeeper;

    private int sessionTimeOut;

    private String lockId;

    public DistributeLock() {
        this.zooKeeper = ZookeeperClient.getInstance();
        this.sessionTimeOut = ZookeeperClient.getSessionTimeOut();
    }

    public static void main(String[] args) {
        CountDownLatch latch = new CountDownLatch(10);
        Random random = new Random();
        for (int i = 0; i < 10; i++) {
            new Thread(() -> {
                DistributeLock lock = null;
                try {
                    lock = new DistributeLock();
                    latch.countDown();
                    latch.await();
                    lock.lock();
                    Thread.sleep(random.nextInt(500));
                } catch (InterruptedException e) {
                    e.printStackTrace();
                } finally {
                    if (null != lock) {
                        lock.unlock();
                    }
                }
            }).start();
        }
    }

    /**
     * 获取锁
     */
    public boolean lock() {

        try {
            //LOCKS/00000001
            lockId = zooKeeper.create(ROOT_LOCKS + "/", data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            System.out.println(Thread.currentThread().getName() + "成功创建了LOCK节点 [" + lockId + "]开始竞争锁");

            //获取根路径下的子节点
            List<String> childrenNodes = zooKeeper.getChildren(ROOT_LOCKS, true);

            //从小到大排序
            SortedSet<String> childrenSet = new TreeSet<String>();
            for (String node : childrenNodes) {
                //因为create成功返回的是全路径
                childrenSet.add(ROOT_LOCKS + "/" + node);
            }
            //获取最小的节点
            String firstNode = childrenSet.first();
            if (lockId.equals(firstNode)) {
                //表示当前就是最小的节点
                System.out.println(Thread.currentThread().getName() + "成功获取锁 [" + lockId + "]");
                return true;
            }

            SortedSet<String> lessThanLockId = childrenSet.headSet(lockId);
            if (!lessThanLockId.isEmpty()) {
                //拿到比当前LOCKID这个节点更小的上一个节点
                String last = lessThanLockId.last();
                zooKeeper.exists(last, event -> {
                    if (event.getType() == Watcher.Event.EventType.NodeDeleted) {
                        countDownLatch.countDown();
                    }
                });
                //如果会话超时或者节点被删除（释放）了
                countDownLatch.await(sessionTimeOut, TimeUnit.MILLISECONDS);
                System.out.println(Thread.currentThread().getName() + "成功获取锁 [" + lockId + "]");
            }

            return true;

        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        return false;
    }

    /**
     * 释放锁
     */
    public boolean unlock() {
        System.out.println(Thread.currentThread().getName() + "开始释放锁[" + lockId + "]");
        try {
            zooKeeper.delete(lockId, -1);
            System.out.println("节点[" + lockId + "]成功被删除");
            return true;
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (KeeperException e) {
            e.printStackTrace();
        }
        return false;
    }

}
